package com.vbobot.demo.spring.cloud.gateway.ribbon.rules.hashing;

import java.util.List;
import java.util.Map.Entry;
import java.util.TreeMap;

/**
 * @author Bobo
 * @date 2022/11/3
 */
public class ConsistentHashLoadBalancer implements LoadBalancer {

    private final HashStrategy hashStrategy;

    private final static int VIRTUAL_NODE_SIZE = 10;
    private final static String VIRTUAL_NODE_SUFFIX = "&&";

    public ConsistentHashLoadBalancer(HashStrategy hashStrategy) {
        this.hashStrategy = hashStrategy;
    }

    @Override
    public <T extends Server> T select(List<T> servers, InvocationObject invoke) {
        final String hashKey = invoke.getHashKey();
        TreeMap<Integer, T> ring = buildConsistentHashRing(servers);
        return select(hashStrategy.getHashCode(hashKey), ring);
    }

    private <T extends Server> T select(int hash, TreeMap<Integer, T> ring) {
        Entry<Integer, T> integerServerEntry = ring.ceilingEntry(hash);
        if (integerServerEntry == null) {
            integerServerEntry = ring.lastEntry();
        }
        return integerServerEntry.getValue();
    }

    private <T extends Server> TreeMap<Integer, T> buildConsistentHashRing(List<T> servers) {
        TreeMap<Integer, T> virtualNodeRing = new TreeMap<>();
        for (T server : servers) {
            for (int i = 0; i < VIRTUAL_NODE_SIZE; i++) {
                virtualNodeRing.put(hashStrategy.getHashCode(server.getUrl() + VIRTUAL_NODE_SUFFIX + i), server);
            }
        }
        return virtualNodeRing;
    }
}
